"""MCP Client Connection Module"""
import os
import asyncio
import logging
import shutil
from contextlib import AsyncExitStack
from typing import Optional, Dict, Any
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

from .tool_manager import ToolManager


class ServerConnection:
    """Handles connection to MCP server with enhanced features"""
    session: Optional[ClientSession] = None
    tool_manager: ToolManager = None

    def __init__(self, name: str, config: Dict[str, Any]):
        """
        Args:
            name: Server name for identification
            config: Server configuration dictionary
        """
        self.name = name
        self.config = config
        self.session: Optional[ClientSession] = None
        self.tool_manager: ToolManager = None
        self.exit_stack = AsyncExitStack()
        self.stdio: Optional[asyncio.StreamReader] = None
        self.write: Optional[asyncio.StreamWriter] = None
        self._cleanup_lock = asyncio.Lock()
        self._retries = config.get("retries", 3)
        self._retry_delay = config.get("retry_delay", 1.0)

    async def connect(self) -> None:
        """Connect to MCP server with retry mechanism"""
        attempt = 0
        while attempt < self._retries:
            try:
                command = (
                    shutil.which(self.config["command"])
                    if self.config["command"] == "npx"
                    else self.config["command"]
                )
                if command is None:
                    raise ValueError("Invalid command specified in config")

                server_params = StdioServerParameters(
                    command=command,
                    args=self.config["args"],
                    env={**os.environ, **self.config["env"]}
                    if self.config.get("env")
                    else None,
                )

                stdio_transport = await self.exit_stack.enter_async_context(
                    stdio_client(server_params)
                )
                self.stdio, self.write = stdio_transport
                self.session = await self.exit_stack.enter_async_context(
                    ClientSession(self.stdio, self.write)
                )
                await self.session.initialize()
                self.tool_manager = ToolManager(server=self.name, session=self.session)
                logging.info(f"Successfully connected to server {self.name}")
                return

            except Exception as e:
                attempt += 1
                logging.warning(
                    f"Connection attempt {attempt} failed for server {self.name}: {e}"
                )
                if attempt < self._retries:
                    await asyncio.sleep(self._retry_delay)
                else:
                    logging.error(f"Max retries reached for server {self.name}")
                    await self.cleanup()
                    raise

    async def cleanup(self) -> None:
        """Clean up server resources safely"""
        async with self._cleanup_lock:
            try:
                if self.session:
                    await self.exit_stack.aclose()
                    self.session = None
                    self.stdio = None
                    self.write = None
                    logging.info(f"Successfully cleaned up server {self.name}")
            except Exception as e:
                logging.error(f"Error during cleanup of server {self.name}: {e}")
                raise

    def is_connected(self) -> bool:
        """Check if connection is active"""
        return self.session is not None

    async def execute_command(self, command: str, arguments: Dict[str, Any]) -> Any:
        """Execute a command on the connected server"""
        if not self.session:
            raise RuntimeError(f"Server {self.name} is not connected")

        try:
            return await self.session.call_tool(command, arguments)
        except Exception as e:
            logging.error(f"Error executing command on server {self.name}: {e}")
            raise
